import pyflag.FlagFramework as FlagFramework
import pyflag.conf
config = pyflag.conf.ConfObject()
import pyflag.FileSystem as FileSystem
import pyflag.Reports as Reports
import pyflag.DB as DB
import os.path
from pyflag.Scanner import *
from pyflag.ColumnTypes import IntegerType, StringType, TimestampType, InodeIDType, FilenameType, ColumnType
import PIL, cStringIO
from pyflag.plugins.DiskForensics.TypeScan import ThumbnailType
from pyflag.FlagFramework import query_type
import pyflag.Reports as Reports
import hashlib
import textwrap
import traceback



class HashMD5Type(ColumnType):
    def __init__(self, **kwargs):
        ColumnType.__init__(self, name="MD5", column="md5", **kwargs)

    def create(self):
        return "`%s` char(32) NOT NULL default ''" % self.column

    def display(self, value, row, result):
        return value.upper()



class HashSHA1Type(ColumnType):
    def __init__(self, **kwargs):
        ColumnType.__init__(self, name="SHA1", column='sha1', **kwargs)

    def create(self):
        return "`%s` char(40) NOT NULL default ''" % self.column

    def display(self, value, row, result):
        return value.upper()
    


class SuspiciousType(StringType):
    def link_display_hook(self, value, row, result):
        if self.link and not self.callback:
            q = self.link.clone()
            q.FillQueryTarget(value)
            tmp = result.__str__()
            result.clear()
            result.link(tmp, q, pane=self.link_pane)
        
    def plain_display_hook(self, value, row, result):
        if value:
            ## Remove non printable chars:
            value = ''.join([ x for x in value if ord(x)>31 ])
            result.text('Click to <br>make suspicious', wrap="full")

    display_hooks = [ plain_display_hook, link_display_hook, ]
    display_hooks_names = [ "plain_display_hook", "link_display_hook", ]
    
    def display(self, value, row, result):
        ## By default just implement a simple callback:
        if self.callback:
            return self.callback(value)
        elif self.wrap:
            value = textwrap.fill( "%s" % value)

        ## Allow all our display hooks to do their things
        for hook in self.display_hooks:
            hook(self, value, row, result)
            


class HashKnownFilesCaseTable(FlagFramework.CaseTable):   
    """ Hash Table - Lists MD5 SHA1 hashes and hash DB Matches """
    name = 'knownfiles'
    columns = [ [ InodeIDType, dict() ],
                [ HashMD5Type, {} ],
                [ HashSHA1Type, {} ],
              ]


class HashUnknownFilesCaseTable(FlagFramework.CaseTable):   
    """ Hash Table - Lists MD5 SHA1 hashes without NSRL Matches """
    name = 'unknownfiles'
    columns = [ [ InodeIDType, dict() ],
                [ HashMD5Type, {} ],
                [ HashSHA1Type, {} ],
              ]


class SuspiciousFilesCaseTable(FlagFramework.CaseTable):
    name = "suspiciousfiles"
    columns = [ [InodeIDType, dict()], 
	      ]	      


class HashScan(GenScanFactory):
    """ Unknown % known files scanner"""
    default = True
    depends = ['TypeScan']
    group = "GeneralForensics"
    
    
    def __init__(self,fsfd):
        GenScanFactory.__init__(self, fsfd)


    class Scan(ScanIfType):
        def __init__(self, inode,ddfs,outer,factories=None,fd=None):
            BaseScanner.__init__(self, inode,ddfs,outer,factories, fd=fd)
            self.m = hashlib.md5()
            self.s= hashlib.sha1()
            self.type = None
            self.length = 0

        def boring(self, metadata, data=''):
            return False
        
        def process(self, data,metadata=None):
            self.boring(metadata,data)
            self.m.update(data)
            self.s.update(data)
            self.length += len(data)

        def finish(self):
            ## Dont do short files
            if self.length<16: return

            dbh_flag=DB.DBO(config.NSRL_DB)
            digestmd5 = self.m.hexdigest()
            digestsha= self.s.hexdigest()

            dbh_flag.execute("select * from hashes where md5='%s' or sha1='%s' limit 1", digestmd5, digestsha)
            pyflaglog.log(8, self.m.hexdigest() + "   " + self.s.hexdigest())
            hashes=dbh_flag.fetch()
            dbh=DB.DBO(self.case)

            dbh.execute("select inode_id from wantedhashes where md5hash='%s' or sha1hash='%s' limit 1", digestmd5, digestsha)
            wanted=dbh.fetch()
            
            inode_id = self.fd.lookup_id()
            if not hashes and not wanted:
                dbh.insert('unknownfiles',
                       inode_id = inode_id,
                       __md5 = digestmd5,
                       __sha1 = digestsha,
                       )
            
            elif not wanted:
                dbh.insert('knownfiles',
                       inode_id = inode_id,
                       __md5 = digestmd5,
                       __sha1 = digestsha,
                       
                       )

class KnownFiles(Reports.report):
    """ Compares MD5 SHA1 hash against the NSRL database to classify files """
    name = "Known Files"
    family = "Hash Analysis"
    description="This report will give a table for describing what the type of file this is based on the MD5 and SHA1 hash matches"
    order = 4

    def form(self,query,result):
        result.case_selector()

    def display(self,query,result):
        result.heading("MD5 SHA1 Hash comparisons")
        dbh=self.DBO(query['case'])
        dbh.check_index('knownfiles','inode_id')
        
        try:
            result.table(
                elements = [ InodeIDType(case=query['case']),
                             FilenameType(case=query['case']),
                             StringType('File Type', 'type', table='type'),
                             HashMD5Type(),
                             HashSHA1Type()],
                table='knownfiles',
                case=query['case'],
                )
        except DB.DBError,e:
            result.para("Error reading the Hash table. Did you remember to run the HashScan scanner?")
            result.para("Error reported was:")
            result.text(e,style="red")



class UnknownFiles(Reports.report):
    name = "Unknown Files"
    family = "Hash Analysis"
    description = "This report will give a table describing files which are not present in NSRL"                               
    order = 3
    
    def form(self,query,result):
        result.case_selector()

    def display(self,query,result):
        self.clear_cache(query)
        result.heading("Unknown files")
        dbh=self.DBO(query['case'])
        dbh.check_index('unknownfiles','inode_id')
	
        try:
            result.table(
                elements = [ SuspiciousType('Suspicious?',column = 'inode_id',table = 'inode', link=query_type(family=query['family'], report='SuspiciousFileForm', case=query['case'], __target__='inode')),
		 	     ThumbnailType(name='Thumbnail', case=query['case']),
                             FilenameType(case=query['case']),
			     StringType('File Type', 'type', table='type'),
                             IntegerType(name='File Size',column='size', table='inode'),
                             HashMD5Type(),
                             HashSHA1Type()],
            	table='unknownfiles',
                case=query['case'])
            
        except DB.DBError,e:
            result.para("Error reading the Hash table. Did you remember to run the HashScan scanner?")
            result.para("Error reported was:")
            result.text(e,style="red")



class SuspiciousFiles(Reports.report):
    name = "Suspicious Files"
    family = "Hash Analysis"
    
    def form(self, query, result):
	result.case_selector()
			
    def display(self, query, result):
	result.heading("Suspicious files")
	dbh = self.DBO(query['case'])
        dbh.expire_cache()
	dbh.check_index('suspiciousfiles', 'inode_id')
							
	try:
	    result.table(
		elements = [ StringType('ID', 'inode_id', table="suspiciousfiles"),
			     ThumbnailType(name='Thumbnail', case = query['case']),	
			     FilenameType(case=query['case']),
			     StringType('File Type', 'type', table = 'type'),
			     IntegerType(name = 'File Size', column = 'size', table = 'inode'),
			     HashMD5Type(table = 'unknownfiles'),
			     HashSHA1Type(table = 'unknownfiles')
			   ],
		table = 'suspiciousfiles',
		case = query['case'])
        except DB.DBError, e:
	    result.para("Error reading the Hash table. Did you remember to run the HashScan scanner?")
	    result.para("Error reported was: ")
	    result.text(e, style = "red")


# Suspicious File Form
class SuspiciousFileForm(Reports.report):
    parameters = {'mark':'string'}
    name = "Suspicious file"
    family = "Hash Analysis"
    hidden = True
    

    def form(self, query, result):
	result.checkbox("Mark file as suspicious", 'mark', 'ok')
    
    def display(self, query, result):
	dbh = self.DBO(query['case'])
	    
	try:
	    if (query['mark'] == 'ok'):
		dbh.execute("insert into suspiciousfiles values(%s)" % query['inode'])
	    result.link("Back", family = "Hash Analysis", report = "Unknown files hashes")
	except DB.DBError, e:
	    result.text(e, style="red")

			     
